from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
import pandas as pd
from typing import Optional, List
from pydantic import BaseModel
import io

app = FastAPI(title="Excel VLOOKUP Service")

class VlookupRequest(BaseModel):
    lookup_column: str  # 查找列名
    return_columns: List[str]  # 需要返回的列
    how: Optional[str] = "left"  # 合并方式

@app.post("/vlookup/")
async def vlookup(
    params: VlookupRequest,  # 必需参数放在前面
    main_file: UploadFile = File(...),  # 主表文件
    lookup_file: UploadFile = File(...)  # 查找表文件
):
    try:
        # 读取上传的Excel文件
        main_df = pd.read_excel(io.BytesIO(await main_file.read()))
        lookup_df = pd.read_excel(io.BytesIO(await lookup_file.read()))

        # 验证列名是否存在
        if params.lookup_column not in main_df.columns or params.lookup_column not in lookup_df.columns:
            raise HTTPException(
                status_code=400,
                detail=f"查找列 {params.lookup_column} 在其中一个文件中不存在"
            )

        for col in params.return_columns:
            if col not in lookup_df.columns:
                raise HTTPException(
                    status_code=400,
                    detail=f"返回列 {col} 在查找表中不存在"
                )

        # 执行merge操作（相当于VLOOKUP）
        result_df = main_df.merge(
            lookup_df[params.return_columns + [params.lookup_column]],
            on=params.lookup_column,
            how=params.how
        )

        # 将结果转换为字典列表
        result = result_df.to_dict(orient='records')

        return JSONResponse(
            content={
                "status": "success",
                "total_rows": len(result),
                "data": result
            }
        )

    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"处理文件时发生错误: {str(e)}"
        )

@app.get("/health")
async def health_check():
    return {"status": "healthy"}
